#  -*- coding: utf-8 -*-

import traceback
import matplotlib.pyplot as plt
import numpy as np
from pymongo import UpdateOne
import pandas as pd
from data.data_module import DataModule
from trading.signal.computer.base_signal_computer import BaseSignalComputer
from util.stock_util import get_all_codes
import mpl_finance as mpf

class BollSignalComputer(BaseSignalComputer):
    def __init__(self):
        BaseSignalComputer.__init__(self, 'boll_signal')

    def compute(self, begin_date, end_date):
        """
        计算指定日期内的信号
        :param begin_date: 开始日期
        :param end_date: 结束日期
        """
        # all_codes = get_all_codes('2018-08-24')
        all_codes=['601398']
        dm = DataModule()

        N = 20
        k = 2
        mistake=0.1
        for code in all_codes[:]:
            print (code)
            try:
                df_daily = dm.get_k_data(code,begin_date=begin_date,autype='kba', end_date=end_date,trade=True)
                # 计算MB，盘后计算，这里用当日的Close
                df_daily['quick'] = df_daily['close'].ewm(span=12).mean()
                # 计算STD20
                df_daily['slow'] = df_daily['close'].ewm(span=26).mean()

                df_daily['diff']= df_daily['quick']-df_daily['slow']

                df_daily['dea']=df_daily['diff'].ewm(span=9).mean()

                df_daily['bar']=2*(df_daily['diff']-df_daily['dea'])
                print (df_daily)
                df_daily.to_csv('erio.csv',encoding='utf-8')
                plt.plot(df_daily['diff'])
                plt.plot(df_daily['dea'])
                df_daily['bar'].plot.bar()
                plt.xticks([])
                plt.show()
            except:
                traceback.print_exc()
            # break

if __name__ == '__main__':
    BollSignalComputer().compute(begin_date='2018-01-01', end_date='2018-08-30')
